import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import tool

class FineLoss(nn.Module):
    def __init__(self, weight=10):
        super(FineLoss, self).__init__()
        self.weight = weight
        self.semi_weight = weight

    def forward(self, ori, option_ori, cor_angle, cor_index):
        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = cor_index

        distance_mac = 2 - 2 * torch.einsum("j,ijk->ik", [ori, option_ori])
        size = distance_mac.shape
        pos = distance_mac[cor_x][cor_y]
        triple_semi_dis = pos - distance_mac[cor_x, :]
        triple_hard_dis1 = pos - distance_mac[:cor_x, :]
        triple_hard_dis2 = pos - distance_mac[cor_x + 1:, :]
        triple_hard_dis = torch.cat([triple_hard_dis1, triple_hard_dis2], dim=0)
        loss = torch.sum(torch.log(1 + torch.exp(triple_hard_dis * self.weight))) + torch.sum(
            torch.log(1 + torch.exp(triple_semi_dis * self.semi_weight)))
        return loss / (size[0] * size[1] - 1)


class CoarseLoss(nn.Module):
    def __init__(self, batch_size, loss_weight):
        super(CoarseLoss, self).__init__()
        self.batch_size = batch_size
        self.loss_weight = loss_weight

    def forward(self, sat_global, grd_global):
        dist_array = 2 - 2 * torch.matmul(sat_global, grd_global.t())
        pos_dist = torch.diag(dist_array)
        pair_n = self.batch_size * (self.batch_size - 1.0)

        # ground to satellite
        triplet_dist_g2s = pos_dist - dist_array
        loss_g2s = torch.sum(torch.log(1 + torch.exp(triplet_dist_g2s * self.loss_weight))) / pair_n

        # satellite to ground
        triplet_dist_s2g = torch.unsqueeze(pos_dist, 1) - dist_array
        loss_s2g = torch.sum(torch.log(1 + torch.exp(triplet_dist_s2g * self.loss_weight))) / pair_n
        loss = (loss_g2s + loss_s2g) / 2.0
        return loss


class TripleLoss2(nn.Module):
    def __init__(self, weight=10):
        super(TripleLoss2, self).__init__()
        self.weight = weight

    def forward(self, ori, option_ori, cor_angle, cor_index):
        option_ori = option_ori.reshape(ori.shape[0], -1, option_ori.shape[1], option_ori.shape[2])
        distance_mac = 2 - 2 * torch.einsum("ik,ijkl->ijl", [ori, option_ori])
        # distance_mac.shape : [baych_size][location][ori]

        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = cor_index
        size = distance_mac.shape
        loss = 0
        for i in range(size[0]):
            pos = distance_mac[i][cor_x[i]][cor_y[i]]
            triple_dis = pos - distance_mac[i]
            loss += torch.sum(torch.log(1 + torch.exp(triple_dis * self.weight)))
            # semi = pos - distance_mac[i][cor_x[i]]
            # loss += torch.sum(torch.log(1 + torch.exp(semi * self.weight)))

        # loss / size[0] / (size[1] * size[2] - 1)
        return loss / size[0]


class TripleLoss(nn.Module):
    def __init__(self, margin=0.4):
        super(TripleLoss, self).__init__()
        self.margin = margin
        self.semi_margin = 1.6

    def forward(self, ori, option_ori, cor_angle, cor_index):

        option_ori = option_ori.reshape(ori.shape[0], -1, option_ori.shape[1], option_ori.shape[2])
        distance_mac = 2 - 2 * torch.einsum("ik,ijkl->ijl", [ori, option_ori])
        # distance_mac.shape : [baych_size][location][ori]

        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = cor_index
        size = distance_mac.shape
        loss = 0
        for i in range(size[0]):
            pos_dist = distance_mac[i][cor_x[i]][cor_y[i]]
            """
            for j in range(size[1]):
                for k in range(size[2]):
                    if j == x:
                        if k != y:
                            loss += self.compute_loss_semi(pos_dist, distance_mac[i][j][k])
                    else:
                        loss += self.compute_loss(pos_dist, distance_mac[i][j][k])
            """
            # fix me: 重写注释部分
            dis = F.relu(- distance_mac[i] + pos_dist + self.margin)
            un_zero_num = torch.count_nonzero(dis) - 1
            loss += (torch.sum(dis) - self.margin) / (un_zero_num + 1e-8)

        return loss / size[0]


# test
if __name__ == '__main__':
    ori = torch.rand(64, 128)
    option_ori = torch.rand(64, 100, 128, 39)

    distance_mac = 2 - 2 * torch.einsum("ik,ijkl->ijl", [ori, option_ori])

    distance_mac_list = []
    for i in range(64):
        ori_ = torch.squeeze(ori[i, :])
        option_ori_ = torch.squeeze(option_ori[i, :, :, :])
        distance = 2 - 2 * torch.einsum("j,ijk->ik", [ori_, option_ori_])
        distance_mac_list.append(distance)

    distance_mac2 = torch.stack(distance_mac_list)
    print(distance_mac2.shape)
    # distance_mac2[0][0][0] *= 1e-5

    print(distance_mac.equal(distance_mac2))